package org.mozhu.simple.rpc.server.network;

import java.io.IOException;
import java.nio.channels.CancelledKeyException;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Selector implements Selectable {
    private java.nio.channels.Selector nioSelector;
    private final Map<String, RpcChannel> channels;
    private final Map<String, RpcChannel> closingChannels;
    private final List<String> failedSends;
    private final int maxReceiveSize;
    private final ChannelBuilder channelBuilder;

    public Selector(ChannelBuilder channelBuilder, int maxReceiveSize) {
        this.channelBuilder = channelBuilder;
        this.maxReceiveSize = maxReceiveSize;
        this.channels = new HashMap<>();
        this.closingChannels = new HashMap<>();
        this.failedSends = new ArrayList<>();
        try {
            this.nioSelector = java.nio.channels.Selector.open();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public void register(String connectionId, SocketChannel socketChannel) throws ClosedChannelException {
        SelectionKey key = socketChannel.register(nioSelector, SelectionKey.OP_READ);
        RpcChannel rpcChannel = channelBuilder.buildChannel(connectionId, key, maxReceiveSize);
        key.attach(rpcChannel);
        this.channels.put(connectionId, rpcChannel);
    }

    /**
     * Return the selector channels.
     */
    public List<RpcChannel> channels() {
        return new ArrayList<>(channels.values());
    }

    /**
     * Return the channel associated with this connection or `null` if there is no channel associated with the
     * connection.
     */
    public RpcChannel channel(String id) {
        return this.channels.get(id);
    }

    /**
     * Return the channel with the specified id if it was disconnected, but not yet closed
     * since there are outstanding messages to be processed.
     */
    public RpcChannel closingChannel(String id) {
        return closingChannels.get(id);
    }

    /**
     * Get the channel associated with selectionKey
     */
    private RpcChannel channel(SelectionKey key) {
        return (RpcChannel) key.attachment();
    }

    public void close(String connectionId) {
        RpcChannel channel = this.channels.get(connectionId);
        if (channel != null) {
            close(channel);
        }
    }

    private void close(RpcChannel channel) {
        channel.disconnect();
    }

    /**
     * Queue the given request for sending in the subsequent {@link #poll(long)} calls
     *
     * @param send The request to send
     */
    public void send(Send send) {
        String connectionId = send.destination();
        if (closingChannels.containsKey(connectionId))
            this.failedSends.add(connectionId);
        else {
            RpcChannel channel = channelOrFail(connectionId, false);
            try {
                channel.setSend(send);
            } catch (CancelledKeyException e) {
                this.failedSends.add(connectionId);
                // close(channel, false);
                close(channel);
            }
        }

    }

    private RpcChannel channelOrFail(String id, boolean maybeClosing) {
        RpcChannel channel = this.channels.get(id);
        if (channel == null && maybeClosing)
            channel = this.closingChannels.get(id);
        if (channel == null)
            throw new IllegalStateException("Attempt to retrieve channel for which there is no connection. Connection id " + id + " existing connections " + channels.keySet());
        return channel;
    }
}
